{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1ce8d73",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import math\n",
    "import torch,random\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "from functools import partial\n",
    "import torchvision.transforms as transforms\n",
    "from matplotlib import pyplot as plt\n",
    "from torch import optim\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import torch.nn.functional as F\n",
    "import torchvision\n",
    "#用于半精度计算\n",
    "from torch.cuda.amp import autocast as autocast\n",
    "#用于ada\n",
    "from diff_augment import DiffAugment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c2200f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#设置一些默认参数\n",
    "dataroot = 'GT'\n",
    "label = 'pre'\n",
    "batch_size = 32\n",
    "image_size = 256\n",
    "ngf = 512  #生成器最终特征大小(生成器的特征容量参数)\n",
    "ndf = 512  #判别器输入特征大小(判别器的特征容量参数)\n",
    "num_epochs = 501  #这个数据要在测试后，正式形成，以便于节约时间\n",
    "lrG = 0.0005   #优化器学习率\n",
    "lrD = 0.0001   #优化器学习率\n",
    "nc = 1        #图像通道数\n",
    "ngpu = 1      #gpu数量\n",
    "workers = 0   #在windows中workers不为0可能会报错\n",
    "nz = 100      #输入向量大小\n",
    "beta1 = 0.5      #原来0.5\n",
    "beta2 = 0.999     #原来0.999\n",
    "capacity_G =  4     #4,3,2,1  这两个参数比较重要，既要节约时间，也要让网络足够稳定，需测试确认\n",
    "capacity_D =  2    #4,3,2,1\n",
    "sa_map_G = image_size//2\n",
    "sa_map_D = image_size//2\n",
    "dpout = True   #是否局部失活，增强泛化性\n",
    "device = torch.device(\"cuda:0\" if (torch.cuda.is_available() and ngpu > 0) else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eba45e3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#制作数据集\n",
    "class MyData(Dataset):\n",
    "    def __init__(self,root_dir,label):\n",
    "        self.root_dir = root_dir\n",
    "        self.label = label\n",
    "        self.path = os.path.join(self.root_dir,self.label)\n",
    "        self.img_path = os.listdir(self.path)\n",
    "    def __getitem__(self, idx):         #一个特殊的函数，__getitem__代表着这个数据集中存储的数据内容的反馈，因此后面可以用dataset[0]调用这里的返回值\n",
    "        img_name = self.img_path[idx]\n",
    "        img_item_path = os.path.join(self.path,img_name)\n",
    "        img = cv2.imread(img_item_path)\n",
    "        # 这个地方可以放入transforms对图像进行tensor化和各种处理,由于G网络中采用tanh激活，因此要将值转换到（-1,1）区间，采用totensor+normalize两步\n",
    "        trans = transforms.Compose([transforms.ToTensor(),transforms.Resize(image_size),transforms.CenterCrop(image_size),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), transforms.Grayscale(num_output_channels=1)])\n",
    "        #trans = transforms.Compose([transforms.ToTensor(), transforms.Grayscale(num_output_channels=1)])\n",
    "        img_tensor = trans(img)\n",
    "        label = self.label\n",
    "        return img_tensor, label\n",
    "    def __len__(self):\n",
    "        return len(self.img_path)\n",
    "dataset = MyData(dataroot,label) #数据集创建成功，类型为np（HW-C）\n",
    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,\n",
    "                                         shuffle=True, num_workers=workers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "862d4e8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "writer = SummaryWriter('logs')\n",
    "step = 0\n",
    "for img in dataloader:\n",
    "    data,target = img\n",
    "    writer.add_images(tag='GPR_data', img_tensor=data, global_step=step, dataformats='NCHW')\n",
    "    step = step + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea613713",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 查看数据集,用plt，但是经过测试plt无法成像tensor数据,因此要转换数据类型\n",
    "def plot_images(img , grid_size = 3):\n",
    "    fig = plt.figure(figsize=(8,8))\n",
    "    columns = rows = grid_size\n",
    "    plt.title('Training Images')\n",
    "    for i in range(1, columns*rows+1):\n",
    "        plt.axis('off')\n",
    "        fig.add_subplot(rows, columns, i)\n",
    "        plt.imshow(dataset[i][0].numpy().squeeze(),cmap='Greys_r')\n",
    "plot_images(dataset) #执行绘图函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a4ad46e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#创建Self-Attention Module，谱归一化\n",
    "class Xiong_Self_Attn(nn.Module):\n",
    "    #适用于图比较大的时候\n",
    "    \"\"\" Self attention Layer\"\"\"\n",
    "    def __init__(self, inChannels, k=4, sn = True):\n",
    "        super(Xiong_Self_Attn,self).__init__()\n",
    "        embedding_channels = inChannels // k  # C_bar\n",
    "        self.sn = sn\n",
    "        if self.sn:\n",
    "            self.key      = nn.utils.spectral_norm(nn.Conv2d(inChannels, embedding_channels, 1))\n",
    "            self.query    = nn.utils.spectral_norm(nn.Conv2d(inChannels, embedding_channels, 1))\n",
    "            self.value    = nn.utils.spectral_norm(nn.Conv2d(inChannels, embedding_channels*2, 1))\n",
    "            self.reprojection = nn.utils.spectral_norm(nn.Conv2d(embedding_channels*2, inChannels, 1))\n",
    "        else:\n",
    "            self.key      = nn.Conv2d(inChannels, embedding_channels, 1)\n",
    "            self.query    = nn.Conv2d(inChannels, embedding_channels, 1)\n",
    "            self.value    = nn.Conv2d(inChannels, embedding_channels*2, 1)\n",
    "            self.reprojection = nn.Conv2d(embedding_channels*2, inChannels, 1)\n",
    "        self.gamma    = nn.Parameter(torch.tensor(0.0), requires_grad=True)      \n",
    "        \n",
    "\n",
    "#        self.softmax  = nn.Softmax(dim=-1) #\n",
    "    def forward(self,x):\n",
    "        \"\"\"\n",
    "            inputs:\n",
    "                x: input feature map [Batch, Channel, Height, Width]\n",
    "            returns:\n",
    "                out: self attention value + input feature\n",
    "                attention: [Batch, Channel, Height, Width]\n",
    "        \"\"\"\n",
    "        batchsize, C, H, W = x.size()\n",
    "        N = H * W                                       # Number of features\n",
    "        f_x = self.key(x).view(batchsize,   -1, N)      # Keys                  [B, C_bar, N]\n",
    "        g_x = self.query(x).view(batchsize, -1, N)      # Queries               [B, C_bar, N]\n",
    "        h_x = self.value(x).view(batchsize, -1, N)      # Values                [B, C_bar, N]\n",
    "        # 对proj_query行执行softmax，对proj_key进行列softmax\n",
    "        g_x = F.softmax(g_x,dim=1)\n",
    "        f_x = F.softmax(f_x,dim=2)\n",
    "        # 计算两次矩阵乘法\n",
    "        G =  torch.bmm(f_x,h_x.permute(0,2,1))#B*dk*dv \n",
    "        v =  torch.bmm(G.permute(0,2,1),g_x )\n",
    "        v = v.view(batchsize, -1, H, W)\n",
    "        o = self.reprojection(v)\n",
    "        out = self.gamma * o + x\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61cef8ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "#创建自适应的网络，保障G和D的网络通过参数轻松设定\n",
    "\n",
    "# 创建生成网络G\n",
    "#第1步，定义GBlock \n",
    "class GBlock(nn.Module):\n",
    "    def __init__(self, in_chan, out_chan, which_bn = nn.BatchNorm2d, activation = nn.ReLU(True), sn = True):\n",
    "        #这里之所以要将各层进行单独命名，是为了后面可能存在的改变各层结构的计划\n",
    "        #默认执行sn\n",
    "        super(GBlock, self).__init__()\n",
    "        \n",
    "        #接收各初始化输入，成为属性（属性表示类除初始化外的函数需要使用）\n",
    "        self.which_bn = which_bn\n",
    "        self.activation = activation\n",
    "        \n",
    "    #定义GBlock中需要用到的各个层\n",
    "        #逆卷积层\n",
    "        if sn:\n",
    "            self.conv = nn.utils.spectral_norm(nn.ConvTranspose2d(in_channels=in_chan, out_channels=out_chan, \n",
    "                                                                  kernel_size=4, stride=2, padding=1, bias=False))\n",
    "        else:\n",
    "            self.conv = nn.ConvTranspose2d(in_channels=in_chan, out_channels=out_chan, \n",
    "                                           kernel_size=4, stride=2, padding=1, bias=False)\n",
    "            \n",
    "        #bn层\n",
    "        self.bn = self.which_bn(out_chan)\n",
    "                   \n",
    "    def  forward(self,x):        \n",
    "        x = self.conv(x)\n",
    "        x = self.bn(x)\n",
    "        x = self.activation(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7796d9a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#第2步，设计生成G网络，主要涉及到：（1）初始网络，（2）注意力网络，（3）最后一层结果\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self, image_size = image_size, latent_dim = nz, last_dim = nc , network_capacity = 4, attn_maps = 256, fmap_max = 512*8, \n",
    "                 snn = True,which_bn = nn.BatchNorm2d, activation = nn.ReLU(True), ngpu = ngpu):   \n",
    "        #网络容量和fmap_max共同限制channels\n",
    "        super(Generator,self).__init__()\n",
    "        #接收各初始化输入，成为属性\n",
    "        self.ngpu = ngpu\n",
    "        self.image_size = image_size\n",
    "        self.latent_dim = latent_dim\n",
    "        self.last_dim = last_dim\n",
    "        self.snn = snn\n",
    "        self.activation = activation\n",
    "        self.which_bn = which_bn\n",
    "        # 定义一些需要的属性        \n",
    "        self.num_layers = int(math.log2(image_size) - 3)  \n",
    "        #需要用到GeneratorBlock的层数，对于生成512*512图来说，需要8,16,32,64,128,256这几层采用GBlock，而1-4,256-512则采用nn.ConvTranspose2d\n",
    "        #最开始的层和最后一层都用简单方法生成,即:1-4,256-512则采用nn.ConvTranspose2d        \n",
    "        filters = [network_capacity * (2 ** (i + 4)) for i in range(self.num_layers)][::-1]#最大的filters = network_capacity*image_size=16*512=4096\n",
    "        set_fmap_max = partial(min, fmap_max)  #最大的filters大小限制，这个原始设置，限定了最大就只有512，所以后面就会是512,512,512,256,128，等等。\n",
    "        filters = list(map(set_fmap_max, filters))  #将所有的大于最大值的都编程最大值，这其实是对通道的一种限制\n",
    "        init_channels = filters[0]*2                  #意味着，4*4的通道数和8*8的通道数一致（这和DCGAN有差异）\n",
    "        filters = [init_channels, *filters]         \n",
    "        #共计7层，将最开始层的大小设计为最大通道数\n",
    "        \n",
    "        in_out_pairs = zip(filters[:-1], filters[1:])  #错位1个，形成output和input通道数的元组\n",
    "        \n",
    "        self.attn_maps = attn_maps\n",
    "        self.attn_layers = int(math.log2(attn_maps) - 3)     #确定SA模块位于一系列残差模块中的位置\n",
    "        \n",
    "    #网络设计部分\n",
    "        #初始层，采用nn.ConvTranspose2d\n",
    "        fist = []\n",
    "        if self.snn:\n",
    "            fist.append(nn.utils.spectral_norm(nn.ConvTranspose2d(in_channels= latent_dim,out_channels= filters[0],\n",
    "                                                                  kernel_size= 4,stride= 1,padding = 0,bias= False)))\n",
    "        else:\n",
    "            fist.append(nn.ConvTranspose2d(in_channels= latent_dim,out_channels= filters[0],\n",
    "                                           kernel_size= 4,stride= 1,padding = 0,bias= False))\n",
    "        fist.append(self.which_bn(filters[0]))\n",
    "        fist.append(self.activation)\n",
    "        self.fist = nn.Sequential(*fist)\n",
    "        \n",
    "        #GBlock+SA模块\n",
    "        self.blocks1 = nn.ModuleList([])\n",
    "        self.blocks2 = nn.ModuleList([])\n",
    "\n",
    "        for ind, (in_chan, out_chan) in enumerate(in_out_pairs):\n",
    "            \n",
    "            if self.attn_layers == ind:             #定义attn层\n",
    "                self.attn = Xiong_Self_Attn(out_chan, sn = self.snn)\n",
    "            \n",
    "            block = GBlock( in_chan, out_chan , sn = self.snn, which_bn = which_bn ,activation = activation)\n",
    "            \n",
    "            if ind > self.attn_layers:      #大于等于SA层则进入blocks2，最极端的情况是不会存在blocks2的\n",
    "                self.blocks2.append(block)\n",
    "            else:\n",
    "                self.blocks1.append(block)   #小于等于SA层则进入blocks1\n",
    "\n",
    "        \n",
    "        #定义最后一层        \n",
    "        last = []\n",
    "        if self.snn:\n",
    "            last.append(nn.utils.spectral_norm(nn.ConvTranspose2d(in_channels= filters[-1], out_channels= self.last_dim, \n",
    "                                                                  kernel_size=4, stride=2, padding=1,bias=False)))\n",
    "        else:\n",
    "            last.append(nn.ConvTranspose2d(in_channels= filters[-1], out_channels= self.last_dim, \n",
    "                                           kernel_size=4, stride=2, padding=1,bias=False))\n",
    "        last.append(nn.Tanh())\n",
    "        self.last = nn.Sequential(*last)\n",
    "        \n",
    "    def forward(self,x):\n",
    "        \n",
    "        x = self.fist(x)    #初始层          \n",
    "                \n",
    "        #一层一层调用 \n",
    "        for m in self.blocks1:     #进行blocks1\n",
    "            x = m(x)\n",
    "            \n",
    "        if self.attn_maps  <  self.image_size:     #当输入512或以上时，就是移除SA模块\n",
    "            x = self.attn(x)                      #判断是否采用SA模块\n",
    "        \n",
    "        if self.attn_maps < self.image_size/2:     # 也就是说，放在256之前，会存在blocks2\n",
    "            for n in self.blocks2:\n",
    "                x = n(x)    \n",
    "                \n",
    "        out = self.last(x)\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a712c9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建判别网络D\n",
    "#第1步，定义DBlock \n",
    "class DBlock(nn.Module):\n",
    "    def __init__(self, in_chan, out_chan, which_bn = nn.BatchNorm2d, activation = nn.LeakyReLU(0.2, inplace=True), sn = True, dpout = False ):\n",
    "        #这里之所以要将各层进行单独命名，是为了后面可能存在的改变各层结构的计划\n",
    "        #默认执行sn\n",
    "        super(DBlock, self).__init__()\n",
    "        \n",
    "        #接收各初始化输入，成为属性（属性表示类除初始化外的函数需要使用）\n",
    "        self.which_bn = which_bn\n",
    "        self.activation = activation\n",
    "        self.dpout = dpout\n",
    "        if self.dpout:\n",
    "          self.dpout_layer = nn.Dropout2d(p=0.2)\n",
    "    #定义DBlock中需要用到的各个层，DBlock和GBlock的唯一区别就是nn.Conv2d和nn.ConvTranspose2d\n",
    "        #卷积层\n",
    "        if sn:\n",
    "            self.conv = nn.utils.spectral_norm(nn.Conv2d(in_channels=in_chan, out_channels=out_chan, \n",
    "                                                         kernel_size=4, stride=2, padding=1, bias=False))\n",
    "        else:\n",
    "            self.conv = nn.Conv2d(in_channels=in_chan, out_channels=out_chan, \n",
    "                                  kernel_size=4, stride=2, padding=1, bias=False)\n",
    "            \n",
    "        #bn层\n",
    "        self.bn = self.which_bn(out_chan)\n",
    "                   \n",
    "    def  forward(self,x):        \n",
    "        x = self.conv(x)\n",
    "        x = self.bn(x)\n",
    "        x = self.activation(x)\n",
    "        if self.dpout:\n",
    "          x = self.dpout_layer(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b83c93d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 第2步，创建判别网络D\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, image_size = image_size, first_dim = nc, network_capacity = 4, attn_maps = 256,  fmap_max = 512*8,\n",
    "                 which_bn = nn.BatchNorm2d, activation = nn.LeakyReLU(0.2, inplace=True), snn = True,ngpu = ngpu,dpout = False):\n",
    "        super(Discriminator,self).__init__()\n",
    "        #接收各初始化输入，成为属性\n",
    "        self.ngpu = ngpu\n",
    "        self.image_size = image_size\n",
    "        self.num_layers = int(math.log2(image_size) - 3) \n",
    "        self.first_dim = first_dim\n",
    "        self.attn_maps = attn_maps\n",
    "        self.snn = snn\n",
    "        self.activation = activation\n",
    "        self.which_bn = which_bn\n",
    "        self.dpout = dpout\n",
    "        #需要用到DBlock的层数，对于生成512*512图来说，需要128，64,32,16,8,4这几层采用DBlock\n",
    "        #最开始的层和最后一层都用简单方法生成即可\n",
    "        \n",
    "        filters = [network_capacity * (2 ** (i + 4)) for i in range(self.num_layers)] #最大的filters = network_capacity*image_size\n",
    "        set_fmap_max = partial(min, fmap_max)       #最大的filters大小限制，这个原始设置，限定了最大就只有512，所以后面就会是512,512,512,256,128，等等。\n",
    "        filters = list(map(set_fmap_max, filters))  #将所有的大于最大值的都编程最大值，这其实是对通道的一种限制\n",
    "        last_channels = filters[-1]*2\n",
    "        filters = [ *filters,last_channels]         \n",
    "        #共计7层，将最开始层的大小设计为最大通道数，第一通道和第二通道的channel数是一样的\n",
    "\n",
    "        in_out_pairs = zip(filters[:-1], filters[1:])  #错位1个，形成output和input通道数的元组\n",
    "                \n",
    "        #网络设计部分\n",
    "        \n",
    "        #定义第1层 ,512-256，相当于G的最后一层\n",
    "        #初始层，采用nn.Conv2d\n",
    "        first = []\n",
    "        if self.snn:\n",
    "            first.append(nn.utils.spectral_norm(nn.Conv2d(in_channels= first_dim,out_channels= filters[0],\n",
    "                                                         kernel_size= 4,stride= 2,padding = 1,bias= False)))\n",
    "        else:\n",
    "            first.append(nn.Conv2d(in_channels= first_dim,out_channels= filters[0],\n",
    "                                  kernel_size= 4,stride= 2,padding = 1,bias= False))\n",
    "        first.append(self.activation)\n",
    "        self.first = nn.Sequential(*first)\n",
    "        \n",
    "        self.attn_layers = int(-(math.log2(attn_maps) - (self.num_layers+1)))  #计算SA模块的位置\n",
    "        \n",
    "        if self.attn_layers == -1:                           #如果为256，则SA模块位于第1层后\n",
    "            self.attn = Xiong_Self_Attn(filters[0],sn = self.snn)\n",
    "\n",
    "        #SA模块和DBlock放在一起\n",
    "        self.blocks1 = nn.ModuleList([])\n",
    "        self.blocks2 = nn.ModuleList([])\n",
    "        \n",
    "        for ind, (in_chan, out_chan) in enumerate(in_out_pairs):\n",
    "            \n",
    "            if self.attn_layers == ind:             #定义attn层\n",
    "                self.attn = Xiong_Self_Attn(out_chan, sn = snn)\n",
    "            \n",
    "            block = DBlock( in_chan, out_chan , sn = self.snn, which_bn = which_bn ,activation = activation, dpout = self.dpout)\n",
    "            if ind <= self.attn_layers:      #小于等于SA层则进入blocks1，极端情况下没有blocks2\n",
    "                self.blocks1.append(block)\n",
    "            else:\n",
    "                self.blocks2.append(block)   #大于进入blocks2\n",
    "\n",
    "        \n",
    "        #定义最后一层,输出为Batchsize,1,1,1        \n",
    "        last = []\n",
    "        if self.snn:\n",
    "            last.append(nn.utils.spectral_norm(nn.Conv2d(in_channels=filters[-1],out_channels=1, kernel_size=4, stride=1, padding=0, bias=False)))\n",
    "        else:\n",
    "            last.append(nn.Conv2d(in_channels=filters[-1], out_channels=1, kernel_size=4, \n",
    "                                  stride=1, padding=0, bias=False))\n",
    "        #从4*4，变到1*1，缩小4倍，暂时不激活\n",
    "        self.last = nn.Sequential(*last)\n",
    "        \n",
    "    def forward(self, x):   #只输入图片就行\n",
    "        x = self.first(x)\n",
    "        if self.attn_layers == -1:      #如果为-1（256），先执行SA，再进入残差网络\n",
    "            x = self.attn(x)\n",
    "            for n in self.blocks2:\n",
    "                x = n(x)\n",
    "        elif self.attn_layers>=0 :                         #大于等于0，先进入blocks1，再进SA，再进blocks2\n",
    "            for m in self.blocks1:\n",
    "                x = m(x)            \n",
    "            if self.attn_maps <  self.image_size/2:  \n",
    "                x = self.attn(x)              \n",
    "                for n in self.blocks2:\n",
    "                    x = n(x) \n",
    "        else :                                            #没有SA\n",
    "            for n in self.blocks2:\n",
    "                x = n(x)\n",
    "                \n",
    "        out = self.last(x)          #最后一层输出\n",
    "           \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e6dd7e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#判别器数据增强（用于ada操作）\n",
    "def random_hflip(tensor, prob):   #随机左右翻转\n",
    "    if prob < random.random():       #小于一个随机数，返回原来的值\n",
    "        return tensor\n",
    "    return torch.flip(tensor, dims=(3,))\n",
    "\n",
    "class AugWrapper(nn.Module):\n",
    "    def __init__(self, D):\n",
    "        super().__init__()\n",
    "        self.D = D\n",
    "    def forward(self, images, prob = 0., types = [], detach = False):\n",
    "        \n",
    "        if random.random() < prob:\n",
    "            images = random_hflip(images, prob=0.5)    #随机左右翻转，概率50%\n",
    "        \n",
    "        if random.random() < prob:\n",
    "            images = DiffAugment(images, types=types)  #随机一系列操作（不是单个操作),这一系列操作本身也具有随机性\n",
    "        if detach:                                      #不用管\n",
    "            images = images.detach()\n",
    "        return self.D(images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45f1f4c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 初始化权重函数，这里可能有问题（经过思考，也许不初始化默认权重比较好）\n",
    "def weights_init(m):\n",
    "    if type(m) == nn.ConvTranspose2d:\n",
    "        nn.init.normal_(m.weight.data,0.0,0.02)\n",
    "    elif type(m) == nn.BatchNorm2d:\n",
    "        nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
    "        nn.init.constant_(m.bias.data, 0)\n",
    "    elif type(m) == nn.Conv2d:\n",
    "        nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
    "    # elif type(m) == nn.InstanceNorm2d:\n",
    "    #     nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
    "    #     nn.init.constant_(m.bias.data, 0)\n",
    "\n",
    "#创建网络实例，传送到多cpu或者cuda，并初始化\n",
    "netG = Generator(image_size = image_size, latent_dim = nz, last_dim = nc , network_capacity = capacity_G, attn_maps = sa_map_G, fmap_max = 512*8, \n",
    "                 snn = True ,which_bn = nn.InstanceNorm2d, activation = nn.ReLU(True), ngpu = ngpu).to(device)\n",
    "netD = Discriminator(image_size = image_size, first_dim = nc, network_capacity = capacity_D, attn_maps = sa_map_D,  fmap_max = 512*8,\n",
    "                     which_bn = nn.InstanceNorm2d, activation = nn.LeakyReLU(0.2, inplace=True) ,snn = True,ngpu = ngpu,dpout = dpout).to(device)\n",
    "netD_aug = AugWrapper(netD).to(device)     #ada网络加入netD\n",
    "\n",
    "if (device.type == 'cuda') and (ngpu > 1):\n",
    "    netG = nn.DataParallel(netG, list(range(ngpu)))\n",
    "if (device.type == 'cuda') and (ngpu > 1):\n",
    "    netD = nn.DataParallel(netD, list(range(ngpu)))\n",
    "\n",
    "netG.apply(weights_init)\n",
    "netD.apply(weights_init)\n",
    "print(netG)             #打印一下两个网络\n",
    "print(netD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "defad4ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "#损失函数和优化器设置\n",
    "criterion = nn.BCEWithLogitsLoss()\n",
    "optimizerG = optim.Adam(netG.parameters(), lr=lrG, betas=(beta1, beta2))\n",
    "optimizerD = optim.Adam(netD.parameters(), lr=lrD, betas=(beta1, beta2))\n",
    "\n",
    "#训练过程中的绘图\n",
    "def gen_img_plot(model,test_input):\n",
    "    predicton = model(test_input).detach().cpu().numpy().squeeze()\n",
    "    fig = plt.figure(figsize=(16,16))     #这个也可以画在tensorboard里面\n",
    "    for i in range(16):\n",
    "        plt.subplot(4,4,i+1)\n",
    "        plt.imshow(predicton[i],cmap='Greys_r')   #(-1,1),听说不能绘图\n",
    "        plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "test_input = torch.randn(16,nz,1,1,device=device)  #这里定义的是16*100的随机数据，能进去G网络？因为G网络要求是1*100输入？？？\n",
    "\n",
    "G_loss = []     #存放每个epoch生成器的损失值\n",
    "D_loss = []     #存放每个epoch判别器的损失值\n",
    "scaler_D = torch.cuda.amp.GradScaler()\n",
    "scaler_G = torch.cuda.amp.GradScaler()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26254941",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard\n",
    "%tensorboard --logdir=logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5a6a2fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "#构建训练循环\n",
    "for epoch in range(num_epochs):\n",
    "    d_epoch_loss = 0    #这里来存计算过程中的每个batch的损失\n",
    "    g_epoch_loss = 0\n",
    "    penalty_epoch_loss = 0\n",
    "    p = 0               #存开始的概率\n",
    "    count = len(dataloader)\n",
    "     #执行每个batch的内部\n",
    "    for setp,(img,_) in enumerate(dataloader):\n",
    "        img = img.to(device)\n",
    "        size = img.size(0)\n",
    "        random_noise = torch.randn(size, nz, 1, 1,device=device)\n",
    "        #构建判别器的训练过程\n",
    "        optimizerD.zero_grad()\n",
    "        \n",
    "        \n",
    "    # 交叉熵损失\n",
    "#         with autocast():\n",
    "#             real_output = netD_aug(img,prob = p ,types = ['translation', 'cutout'])   #判别器输入真实图片，得到真实图片的预测real_output\n",
    "#             d_real_loss = criterion(real_output,torch.ones_like(real_output))    #G网络中真实图片损失\n",
    "\n",
    "\n",
    "#             gen_img = netG(random_noise)\n",
    "#             fake_output = netD_aug(gen_img.detach(),prob = p ,types = ['translation', 'cutout'],detach = True) \n",
    "#             d_fake_loss = criterion(fake_output,torch.zeros_like(fake_output))    #G网络中假图片损失\n",
    "#             d_loss = d_real_loss + d_fake_loss\n",
    "            \n",
    "#         scaler_D.scale(d_loss).backward()\n",
    "#         scaler_D.step(optimizerD)\n",
    "#         scaler_D.update()\n",
    "\n",
    "#         #构建生成器的训练过程\n",
    "#         optimizerG.zero_grad()\n",
    "#         with autocast():\n",
    "#           fake_output = netD_aug(gen_img,prob = p ,types = ['translation', 'cutout'])\n",
    "#           g_loss = criterion(fake_output,torch.ones_like(fake_output))\n",
    "#         scaler_G.scale(g_loss).backward()\n",
    "#         scaler_G.step(optimizerG)\n",
    "#         scaler_G.update()\n",
    "# #     #softpluse 损失\n",
    "#         with autocast():\n",
    "#             real_output = netD_aug(img,prob = p ,types = ['translation', 'cutout'])   #判别器输入真实图片，得到真实图片的预测real_output\n",
    "#             d_real_loss =  F.softplus(-real_output).mean()\n",
    "        \n",
    "#             gen_img = netG(random_noise)\n",
    "#             fake_output = netD_aug(gen_img.detach(),prob = p ,types = ['translation', 'cutout'],detach = True) \n",
    "#             d_fake_loss = F.softplus(fake_output).mean()\n",
    "#             d_loss =  d_real_loss + d_fake_loss\n",
    "        \n",
    "#         scaler_D.scale(d_loss).backward()\n",
    "#         scaler_D.step(optimizerD)\n",
    "#         scaler_D.update()\n",
    "\n",
    "#         #构建生成器的训练过程\n",
    "#         optimizerG.zero_grad()\n",
    "#         with autocast():\n",
    "#             fake_output = netD_aug(gen_img,prob = p ,types = ['translation', 'cutout'])\n",
    "#             g_loss = F.softplus(-fake_output).mean()\n",
    "#         scaler_G.scale(g_loss).backward()\n",
    "#         scaler_G.step(optimizerG)\n",
    "#         scaler_G.update()\n",
    "\n",
    "    #hinge损失        \n",
    "        with autocast():\n",
    "#             real_output = netD(img)   #判别器输入真实图片，得到真实图片的预测real_output\n",
    "            real_output = netD_aug(img,prob = p ,types = ['translation', 'cutout'])   #判别器输入真实图片，得到真实图片的预测real_output\n",
    "            d_real_loss =  F.relu(1.-real_output).mean()\n",
    "        \n",
    "            gen_img = netG(random_noise)\n",
    "#             fake_output = netD(gen_img.detach())   #判别器输入假图片，得到假图片的预测fake_output.注意由于训练判别器的过程中要保持生成器不变，所以要截断梯度\n",
    "            fake_output = netD_aug(gen_img.detach(),prob = p ,types = ['translation', 'cutout'],detach = True) \n",
    "            #img.clone().detach()，通过d_epoch_loss控制概率\n",
    "            d_fake_loss = F.relu(1.+fake_output).mean()\n",
    "            d_loss =  d_real_loss + d_fake_loss\n",
    "        \n",
    "        scaler_D.scale(d_loss).backward()\n",
    "        scaler_D.step(optimizerD)\n",
    "        scaler_D.update()\n",
    "\n",
    "        #构建生成器的训练过程\n",
    "        optimizerG.zero_grad()\n",
    "        with autocast():\n",
    "#             fake_output = netD(gen_img)   #生成假图 \n",
    "            fake_output = netD_aug(gen_img,prob = p ,types = ['translation', 'cutout'])   \n",
    "            g_loss = -torch.mean(fake_output)\n",
    "        scaler_G.scale(g_loss).backward()\n",
    "        scaler_G.step(optimizerG)\n",
    "        scaler_G.update() \n",
    "                            \n",
    "        print(\"Epoch:\", epoch,\"Batch:\",setp,\"/\",count)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            d_epoch_loss += d_loss\n",
    "            g_epoch_loss += g_loss\n",
    "    with torch.no_grad():\n",
    "        d_epoch_loss /= count\n",
    "        g_epoch_loss /= count\n",
    "        #加入ada的概率控制\n",
    "        if 0 <= d_epoch_loss < 1.0:    \n",
    "            p = 1-d_epoch_loss\n",
    "            p = min(0.8,p)\n",
    "            \n",
    "        G_loss.append(d_epoch_loss)\n",
    "        D_loss.append(d_epoch_loss)\n",
    "        print(\"Epoch:\",epoch)\n",
    "        writer.add_scalars('loss', {\"Gloss\":g_epoch_loss,'Dloss':d_epoch_loss}, global_step=epoch) \n",
    "        #将损失输入tensorbo\n",
    "    if epoch % 20 == 0:           #看训练过程中的图片\n",
    "        gen_img_plot(netG, test_input)\n",
    "    if epoch % 50 == 0:           #看训练过程中的图片\n",
    "        torch.save(netG.state_dict(),\"netG42{}.pth\".format(epoch))    #既保存了结构，也保存了参数\n",
    "writer.close()\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "f1=open(\"G_loss.txt\",\"w\")     #保存至txt中，每次运行都会覆盖原来txt中的内容\n",
    "for i in range(len(G_loss)):\n",
    "    f1.write(str(G_loss[i]))\n",
    "    f1.write(\"\\n\")\n",
    "f1.close()\n",
    "f2=open(\"D_loss.txt\",\"w\")\n",
    "for i in range(len(D_loss)):\n",
    "    f2.write(str(D_loss[i]))\n",
    "    f2.write(\"\\n\")\n",
    "f2.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a096b84",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The following code is used for the inference stage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "959beba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#验证模型\n",
    "#假设生成100张图\n",
    "\n",
    "#第一步，定义模型,并加载参数\n",
    "netG = Generator(image_size = image_size, latent_dim = nz, last_dim = nc , network_capacity = capacity_G, attn_maps = sa_map_G, fmap_max = 512*8, \n",
    "                 snn = True,which_bn = nn.InstanceNorm2d, activation = nn.ReLU(True), ngpu = ngpu).to(device)\n",
    "name = '512netG42-M150.pth'\n",
    "# print(netG)\n",
    "netG.load_state_dict(torch.load(name,map_location=torch.device('cpu')))\n",
    "if os.path.exists('{}_img'.format(name)):\n",
    "    pass\n",
    "else:\n",
    "    os.makedirs('{}_img'.format(name))  #创建文件夹\n",
    "\n",
    "# #第二步，定义随机变量\n",
    "test_input = torch.randn(50,nz,1,1,device=device)  #这里定义的是16*100的随机数据，能进去G网络？因为G网络要求是1*100输入？？？\n",
    "\n",
    "# #第三步，验证\n",
    "netG.eval()             #开始测试\n",
    "for i in range(test_input.size(0)):\n",
    "    with torch.no_grad():    #测试的时候梯度不需要改变 参数不需要更新 要保持良好的代码习惯哦~\n",
    "        output = netG(test_input[i].unsqueeze(0))\n",
    "        output_img = (output+1)/2   #变换到0-1空间中\n",
    "#     print(\"Number:\",i)\n",
    "        #保存图片\n",
    "    torchvision.utils.save_image(output_img.squeeze(),\"{}_img\\{}_{}.jpg\".format(name,name,i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b97992c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#验证模型\n",
    "#假设生成100张图\n",
    "\n",
    "#第一步，定义模型,并加载参数\n",
    "netG = Generator(image_size = image_size, latent_dim = nz, last_dim = nc , network_capacity = capacity_G, attn_maps = sa_map_G, fmap_max = 512*8, \n",
    "                 snn = True,which_bn = nn.InstanceNorm2d, activation = nn.ReLU(True), ngpu = ngpu).to(device)\n",
    "model_name = os.listdir('model')\n",
    "for j in range(len(model_name)):\n",
    "    name = os.path.join('model',model_name[j])\n",
    "    netG.load_state_dict(torch.load(name,map_location=torch.device('cpu')))\n",
    "    if os.path.exists('{}_img'.format(model_name[j])):\n",
    "        pass\n",
    "    else:\n",
    "        os.makedirs('{}_img'.format(model_name[j]))  #创建文件夹\n",
    "\n",
    "    # #第二步，定义随机变量\n",
    "    test_input = torch.randn(50,nz,1,1,device=device)  #这里定义的是16*100的随机数据，能进去G网络？因为G网络要求是1*100输入？？？\n",
    "\n",
    "    # #第三步，验证\n",
    "    netG.eval()             #开始测试\n",
    "    for i in range(test_input.size(0)):\n",
    "        with torch.no_grad():    #测试的时候梯度不需要改变 参数不需要更新 要保持良好的代码习惯哦~\n",
    "            output = netG(test_input[i].unsqueeze(0))\n",
    "            output_img = (output+1)/2   #变换到0-1空间中\n",
    "#     print(\"Number:\",i)\n",
    "        #保存图片\n",
    "        torchvision.utils.save_image(output_img.squeeze(),\"{}_img\\{}_{}.jpg\".format(model_name[j],model_name[j],i))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "'YOLOV5'",
   "language": "python",
   "name": "yolov5"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
